from geesibling.core.types import Graph, Node
from geesibling.core.lib._graph import DataType,DeviceType,Device,search_policy
from geesibling.tools import log
from geesibling.adapters.pytorch.getTorchGraph import is_tensor,getTorchGraph

from torchvision import models
import torch
from torch import fx

__doc__ = """
convert torch.fx graph to Geesibling Graph
Author: jirongting
datetime: 2024.04.15
version: 2 2024.04.15 second commit
"""

"""
add inputs for each geesibling node

Parameters:
 node - geesibling Node
 fx_node - torch.fx Node

"""
def add_inputs(node:Node, fx_node:fx.Node):
    for input_node in fx_node._input_nodes:
        node.add_input(input_node.name)


"""
add outputs for each geesibling node

Parameters:
 node - geesibling Node
 fx_node - torch.fx Node

"""
def add_outputs(node:Node, fx_node:fx.Node):
    for output_node in fx_node.users:
        node.add_output(output_node.name)

"""
add input_port for each geesibling node

Parameters:
 node - geesibling Node
 fx_node - torch.fx Node

"""

def add_abstract_inputs(node:Node,fx_node:fx.Node):
    for i, input_node in enumerate(fx_node._input_nodes):
        # AddInputPort(input_node, input_index, index, dtype, shape),input_index??
        if 'val' in input_node.meta: # if the input node has output, add the shape into inputport else []:
            if isinstance(input_node.meta['val'],torch.Tensor):
                node.add_inputport(input_node.name,0,i,DataType.F32,list(input_node.meta['val'].data.shape))
            else:
                val = input_node.meta['val'][fx_node.args[1]]
                node.add_inputport(input_node.name,0,i,DataType.F32,list(val.data.shape))
        else:
            node.add_inputport(input_node.name,0,i,DataType.F32,[])

"""
add output_port for each geesibling node

Parameters:
 node - geesibling Node
 fx_node - torch.fx Node

"""
def add_abstract_outputs(node:Node,fx_node:fx.Node):

    if 'val' in fx_node.meta: # if the input node has output, add the shape into inputport else []:
        if isinstance(fx_node.meta['val'],torch.Tensor):
            node.add_outputport(DataType.F32,list(fx_node.meta['val'].data.shape),0)
        else:
            if fx_node.meta['val'] and not isinstance(fx_node.meta['val'], bool):
                for i, val in enumerate(fx_node.meta['val']):
                    node.add_outputport(DataType.F32,list(val.data.shape),i)
            else:
                node.add_outputport(DataType.F32,[],0)
    else:
        node.add_outputport(DataType.F32,[],0)
"""
fx graph convert to Geesibling Graph

Parameters:
 traced - transformers.utils.fx graph

Returns:
 return a Geesibling Graph

"""
def add_attrs(node:Node,fx_node:fx.Node):
    attrs={}
    if 'stack_trace' in fx_node.meta:   #是参数
        attrs['target'] = str(fx_node.target)
    else:   #是输入
        attrs['target'] = "USER_INPUT"
    for i, arg in enumerate(fx_node.args):
        if not isinstance(arg,torch.fx.Node):
            attrs['immutable_list'] = str(arg)
    node.attrs = attrs
    # if 'stack_trace' in fx_node.meta:   #是参数
    #     node.attrs = {"target":str(fx_node.target)}
    # else:   #是输入
    #     node.attrs = {"target":"USER_INPUT"}

def graphToGeeGraph(traced:fx.graph_module.GraphModule):
    geesiblingGraph = Graph()
    print("graph")
    print(traced)
    for fx_node in traced.graph.nodes:
        node = Node(fx_node.name,fx_node.op) # attr name,op
        add_inputs(node,fx_node) # attr inputs
        add_outputs(node,fx_node)# attr outputs
        add_abstract_inputs(node,fx_node)
        add_abstract_outputs(node,fx_node)
        add_attrs(node,fx_node)
        geesiblingGraph.add_node(node)
    return geesiblingGraph
